import os
import argparse
import pickle
import itertools
import numpy as np
import scipy.spatial
from tqdm import tqdm


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('-i', '--input', type=str, help='path of centers', required=True)
    parser.add_argument('-o', '--output', type=str, help='path of output map', required=True)
    args = parser.parse_args()
    return args


if __name__ == '__main__':
    args = parse_args()

    with open(args.input, 'rb') as f:
        centers = pickle.load(f)

    kdtree = scipy.spatial.cKDTree(centers)

    nearest_map = np.empty((256, 256, 256), dtype=np.uint8)
    for r, g, b in tqdm(itertools.product(range(256), range(256), range(256))):
        nearest_map[r, g, b] = kdtree.query(np.array([[r, g, b]]), k=1)[1][0]

    os.makedirs(os.path.dirname(args.output), exist_ok=True)
    nearest_map.dump(args.output)

    print("done")
